from sparkai.llm.llm import ChatSparkLLM, ChunkPrintHandler
from sparkai.core.messages import ChatMessage

# 星火认知大模型的URL
SPARKAI_URL = 'wss://spark-api.xf-yun.com/v1.1/chat'
# 在控制台获取的API密钥
SPARKAI_APP_ID = '6e4244ba'
SPARKAI_API_SECRET = 'NGJmNjAyNDcxYzdmNGJhMzI2NTFhMWI4'
SPARKAI_API_KEY = 'a70a751b89b2a2174fed3fd1328cdaa3'
# 使用的模型版本
SPARKAI_DOMAIN = 'general'

if __name__ == '__main__':
    spark = ChatSparkLLM(
        spark_api_url=SPARKAI_URL,
        spark_app_id=SPARKAI_APP_ID,
        spark_api_key=SPARKAI_API_KEY,
        spark_api_secret=SPARKAI_API_SECRET,
        spark_llm_domain=SPARKAI_DOMAIN,
        streaming=False,  # 设置为True即可启用流式返回
    )

    messages = [ChatMessage(
        role="user",
        content='你是谁'
    )]
    
    handler = ChunkPrintHandler()
    response = spark.generate([messages], callbacks=[handler])
    print(response)